from preprocess.file_paser.docx_paser.docx_utils import get_content
from preprocess.data_utils import traverse_dir
from docx import Document
import os
import json
from tqdm import tqdm


def _parse_docx_content(file_type, file_dir, image_saved_dir):
    if not os.path.exists(image_saved_dir):
        os.makedirs(image_saved_dir)
    iter_docx = traverse_dir(file_dir, ['docx'])
    docx_contents = []
    progress_bar = enumerate(tqdm(iter_docx, desc=f'正在提取{file_type}文件内容'))
    for fid, (file_path, file_name) in progress_bar:
        fid = f'{file_type}_{fid}'
        doc = Document(file_path)
        doc_content, image_info, table_info = get_content(fid, doc, img_save_path=image_saved_dir)
        item = {'file_name': file_name, 'image_info': image_info, 'table_info': table_info, 'content': doc_content, }
        docx_contents.append(item)
    return docx_contents


def parse_docx(docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'docx'
    docx_json_path = f'{out_dir}/docx_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(file_suffix, docx_file_dir, image_saved_dir)
    with open(docx_json_path, 'w+', encoding='utf-8') as f:
        data = {'file_num': len(docx_contents), 'file_suffix': '.docx', 'file_list': docx_contents}
        json.dump(data, f, ensure_ascii=False)


def parse_pdf(docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'pdf'
    docx_json_path = f'{out_dir}/pdf_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(file_suffix, docx_file_dir, image_saved_dir)
    with open(docx_json_path, 'w+', encoding='utf-8') as f:
        data = {'file_num': len(docx_contents), 'file_suffix': '.pdf', 'file_list': docx_contents}
        json.dump(data, f, ensure_ascii=False)


def parse_doc(docx_file_dir, out_dir, image_saved_dir):
    file_suffix = 'doc'
    docx_json_path = f'{out_dir}/doc_content.json'
    if os.path.exists(docx_json_path):
        os.remove(docx_json_path)
    docx_contents = _parse_docx_content(file_suffix, docx_file_dir, image_saved_dir)
    with open(docx_json_path, 'w+', encoding='utf-8') as f:
        data = {'file_num': len(docx_contents), 'file_suffix': '.doc', 'file_list': docx_contents}
        json.dump(data, f, ensure_ascii=False)
